import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
import mysql.connector
from datetime import datetime
import joblib # 用于保存和加载模型

# ----------------------------
# 解决中文显示问题（Windows系统）
# ----------------------------
plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei"]
plt.rcParams['axes.unicode_minus'] = False

# ----------------------------
# 数据库配置（请根据您的MySQL配置修改）
# ----------------------------
DB_CONFIG = {
    'host': 'localhost',
    'user': 'root',     # <--- 替换为您的 MySQL 用户名
    'password': '0000', # <--- 替换为您的 MySQL 密码
    'database': 'product_management_db' # 您的数据库名称
}

# ----------------------------
# 辅助函数：数据库操作
# ----------------------------
def get_db_connection():
    """获取数据库连接"""
    try:
        conn = mysql.connector.connect(**DB_CONFIG)
        return conn
    except mysql.connector.Error as err:
        print(f"数据库连接失败: {err}")
        return None

def create_tables_mysql():
    """创建 MySQL 数据库表"""
    conn = get_db_connection()
    if conn is None:
        return
    try:
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS predictions (
                id INT AUTO_INCREMENT PRIMARY KEY,
                original_index INT,
                actual_value DOUBLE,
                custom_pred DOUBLE,
                sklearn_pred DOUBLE,
                model_name VARCHAR(255),
                run_date DATETIME
            )
        ''')
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS model_metrics (
                id INT AUTO_INCREMENT PRIMARY KEY,
                model_type VARCHAR(50),
                metric_name VARCHAR(50),
                metric_value DOUBLE,
                run_date DATETIME
            )
        ''')
        conn.commit()
        print(f"\nMySQL 数据库 '{DB_CONFIG['database']}' 中的表已创建/检查。")
    except mysql.connector.Error as err:
        print(f"创建表时发生 MySQL 错误: {err}")
    finally:
        if conn and conn.is_connected():
            cursor.close()
            conn.close()

def insert_prediction_data_mysql(y_test_series, y_pred_custom_array, y_pred_sklearn_array):
    """插入预测结果数据到 MySQL"""
    conn = get_db_connection()
    if conn is None:
        return
    try:
        cursor = conn.cursor()
        run_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        if not isinstance(y_test_series, pd.Series):
            y_test_series = pd.Series(y_test_series) # 尝试转换为Series
            print("Warning: y_test_series was not a Pandas Series, converted for index handling.")

        data_to_insert = []
        for i, original_idx in enumerate(y_test_series.index):
            actual = float(y_test_series.iloc[i])
            custom_p = float(y_pred_custom_array[i]) if i < len(y_pred_custom_array) else None
            sklearn_p = float(y_pred_sklearn_array[i]) if i < len(y_pred_sklearn_array) else None
            
            data_to_insert.append((
                int(original_idx), actual, custom_p, sklearn_p, "Comparison Run", run_date
            ))

        sql = '''
            INSERT INTO predictions (original_index, actual_value, custom_pred, sklearn_pred, model_name, run_date)
            VALUES (%s, %s, %s, %s, %s, %s)
        '''
        cursor.executemany(sql, data_to_insert)
        
        conn.commit()
        print(f"成功插入 {len(data_to_insert)} 条预测数据到 '{DB_CONFIG['database']}.predictions'。")
    except mysql.connector.Error as err:
        print(f"插入预测数据时发生 MySQL 错误: {err}")
    except Exception as e:
        print(f"插入预测数据时发生其他错误: {e}")
    finally:
        if conn and conn.is_connected():
            cursor.close()
            conn.close()

def insert_metric_data_mysql(model_type, metric_name, metric_value):
    """插入模型评估指标数据到 MySQL"""
    conn = get_db_connection()
    if conn is None:
        return
    try:
        cursor = conn.cursor()
        run_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        sql = '''
            INSERT INTO model_metrics (model_type, metric_name, metric_value, run_date)
            VALUES (%s, %s, %s, %s)
        '''
        cursor.execute(sql, (model_type, metric_name, float(metric_value), run_date))
        
        conn.commit()
        print(f"成功插入指标：{model_type} - {metric_name}: {metric_value:.4f}")
    except mysql.connector.Error as err:
        print(f"插入指标数据时发生 MySQL 错误: {err}")
    except Exception as e:
        print(f"插入指标数据时发生其他错误: {e}")
    finally:
        if conn and conn.is_connected():
            cursor.close()
            conn.close()

# ----------------------------
# 1. 加载数据集
# ----------------------------
try:
    data_path = os.path.join(os.path.dirname(__file__), 'data', 'new_data.csv')
    df = pd.read_csv(data_path)
    print("数据集加载成功！")
except FileNotFoundError:
    print(f"错误：未找到文件 {data_path}")
    exit()
except Exception as e:
    print(f"加载数据时发生错误：{str(e)}")
    exit()

# ----------------------------
# 2. 数据预处理
# ----------------------------
target_column = 'value'
if target_column not in df.columns:
    print(f"错误：数据集中未找到目标列'{target_column}'")
    exit()

X = df.drop(columns=[target_column])
y = df[target_column]

# 记录原始特征列名，用于后续API预测时对齐
original_features = X.columns.tolist()

non_numeric_features = X.select_dtypes(exclude=['number']).columns
if len(non_numeric_features) > 0:
    print(f"编码非数值特征：{non_numeric_features.tolist()}")
    X = pd.get_dummies(X, columns=non_numeric_features)

if X.isnull().sum().sum() > 0:
    X = X.fillna(X.mean())
    y = y.fillna(y.mean())

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42
)

# 保存训练集的列名，以便后续API预测时对齐输入特征
training_features = X.columns.tolist()
joblib.dump(training_features, os.path.join(os.path.dirname(__file__), 'training_features.pkl'))
print(f"训练特征列名已保存到 training_features.pkl. 数量: {len(training_features)}")


# ----------------------------
# 3. 自定义随机森林回归器（保持不变）
# ----------------------------
class RandomForestRegressorCustom:
    """自定义随机森林回归器"""
    def __init__(self, n_estimators=100, random_state=0):
        self.n_estimators = n_estimators
        self.random_state = random_state
        self.trees = []
        self.feature_names = None
    
    def fit(self, X, y):
        self.feature_names = X.columns.tolist() if isinstance(X, pd.DataFrame) else [f'feature_{i}' for i in range(X.shape[1])]
        n_samples = X.shape[0]
        rs = np.random.RandomState(self.random_state)
        
        for _ in range(self.n_estimators):
            dt = DecisionTreeRegressor(
                random_state=rs.randint(np.iinfo(np.int32).max),
                max_features="sqrt"
            )
            sample_indices = rs.randint(0, n_samples, n_samples)
            sample_weight = np.bincount(sample_indices, minlength=n_samples)
            dt.fit(X, y, sample_weight=sample_weight)
            self.trees.append(dt)
    
    def predict(self, X):
        if isinstance(X, pd.DataFrame):
            # 确保输入X的列与训练时的特征列对齐，缺失的补0
            X_aligned = X.reindex(columns=self.feature_names, fill_value=0)
        else: # 如果是非DataFrame输入，尝试转换为DataFrame
             try:
                 X_aligned = pd.DataFrame([X], columns=self.feature_names) # 假设X是单个样本的列表/数组
             except ValueError:
                 print("Error: Input X for prediction is not a DataFrame and cannot be converted with provided features.")
                 return np.array([]) # 返回空数组或错误
        
        y_pred = np.zeros(X_aligned.shape[0])
        for dt in self.trees:
            y_pred += dt.predict(X_aligned)
        
        return y_pred / self.n_estimators

# ----------------------------
# 4. 模型训练与评估
# ----------------------------
print("\n开始训练自定义随机森林模型...")
rf_reg_custom = RandomForestRegressorCustom(n_estimators=100, random_state=42)
rf_reg_custom.fit(X_train, y_train)

# 预测
y_pred_custom = rf_reg_custom.predict(X_test)

# 评估
mse_custom = mean_squared_error(y_test, y_pred_custom)
r2_custom = r2_score(y_test, y_pred_custom)
print("\n===== 自定义随机森林回归评估 =====")
print(f"均方误差 (MSE): {mse_custom:.4f}")
print(f"决定系数 (R²): {r2_custom:.4f}")

# 保存自定义模型
model_path_custom = os.path.join(os.path.dirname(__file__), 'trained_model_custom_rf.pkl')
joblib.dump(rf_reg_custom, model_path_custom)
print(f"自定义随机森林模型已保存到 {model_path_custom}")

print("\n开始训练 Scikit-learn 随机森林模型...")
rf_sklearn = RandomForestRegressor(n_estimators=100, random_state=42)
rf_sklearn.fit(X_train, y_train)
y_pred_sklearn = rf_sklearn.predict(X_test)

mse_sklearn = mean_squared_error(y_test, y_pred_sklearn)
r2_sklearn = r2_score(y_test, y_pred_sklearn)
print("\n===== Scikit-learn 随机森林回归评估 =====")
print(f"均方误差 (MSE): {mse_sklearn:.4f}")
print(f"决定系数 (R²): {r2_sklearn:.4f}")

# 保存 Scikit-learn 模型
model_path_sklearn = os.path.join(os.path.dirname(__file__), 'trained_model_sklearn_rf.pkl')
joblib.dump(rf_sklearn, model_path_sklearn)
print(f"Scikit-learn 随机森林模型已保存到 {model_path_sklearn}")


# ----------------------------
# 6. 将评估指标存入 MySQL 数据库
# ----------------------------
create_tables_mysql() # 确保表存在

# 存储自定义模型指标
insert_metric_data_mysql("custom_rf", "MSE", mse_custom)
insert_metric_data_mysql("custom_rf", "R2", r2_custom)

# 存储 sklearn 模型指标
insert_metric_data_mysql("sklearn_rf", "MSE", mse_sklearn)
insert_metric_data_mysql("sklearn_rf", "R2", r2_sklearn)

# 存储训练时的预测结果（可选，如果需要将测试集的预测也保存）
# insert_prediction_data_mysql(y_test, y_pred_custom, y_pred_sklearn)

print("\n数据存储到 MySQL 数据库完成。")

# 可视化部分（保持不变，但通常在生产环境中不显示）
if X_train.shape[1] > 0:
    # 可以在这里选择不显示图形或者将图形保存到文件
    print("\n可视化图表已生成（如果需要显示，请取消注释 plt.show()）")
    # plt.show() # 在生产环境中通常不显示，可以注释掉或保存为图片文件